/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "custom_all_reduce_operator.h"
#include "coll_alg_op_registry.h"

namespace hccl {

CustomAllReduceOperator::CustomAllReduceOperator(AlgConfigurator *algConfigurator, CCLBufferManager &cclBufferManager,
                                                 HcclDispatcher dispatcher, std::unique_ptr<TopoMatcher> &topoMatcher)
    : CollAlgOperator(algConfigurator, cclBufferManager, dispatcher, topoMatcher, HcclCMDType::HCCL_CMD_ALLREDUCE)
{
}

CustomAllReduceOperator::~CustomAllReduceOperator() {}

HcclResult CustomAllReduceOperator::SelectAlg(const std::string &tag, const OpParam &param, std::string &algName,
                                              std::string &newTag)
{
    constexpr u64 HCCL_CONTEST_SMALL_COUNT = 512 * 1024;       // 512KB
    constexpr u64 HCCL_CONTEST_MEDIUM_COUNT = 2 * 1024 * 1024; // 2MB
    constexpr u64 HCCL_CONTEST_HUGE_COUNT = 64 * 1024 * 1024;  // 64MB
    algName = "CustomSmallAllReduceMeshExecutor";
    return HCCL_SUCCESS;
    // 算法选择逻辑
    // TODO: 选手可根据数据量大小选择合适的 Executor
    // 注意:
    // 1. 相同算法在不同数据量下的性能不同
    // 2. 选手可以先只实现一个 Executor，算法选择时直接设置 algName 为该 Executor 的名字

    u32 unitSize = SIZE_TABLE[param.DataDes.dataType];
    u64 dataSize = param.DataDes.count * unitSize; // 单位：字节，三种数据量：512K/2M/64M
    if (dataSize <= HCCL_CONTEST_SMALL_COUNT) {
        algName = "CustomSmallAllReduceMeshExecutor";
    } else if (dataSize <= HCCL_CONTEST_MEDIUM_COUNT) {
        algName = "CustomMediumAllReduceMeshExecutor";
    } else {
        algName = "CustomHugeAllReduceMeshExecutor";
    }
    return HCCL_SUCCESS;
}

// 注册算子
REGISTER_OP(HcclCMDType::HCCL_CMD_ALLREDUCE, AllReduce, CustomAllReduceOperator);
} // namespace hccl
